package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

import nlp.langmodel.LanguageModel;
import nlp.util.Counter;
import nlp.util.CounterMap;

public class BigramModel implements LanguageModel {
	private Collection<List<String>> trainingCollection;
	private Collection<List<String>> validationCollection;
//	private Counter<String> wordCounter = new Counter<String>();
	private CounterMap<String, String> bigramCounterMap = new CounterMap<String, String>();
	private Counter<Integer> nCounter = new Counter<Integer>();
	private Counter<String> historyCounter = new Counter<String>();
	private static final double INCREMENT_FACTOR = 1.0;
	private final static Integer THRESHOLD = 5;
	private Double[] N = new Double[THRESHOLD + 2];
	private Double[] C = new Double[THRESHOLD + 1];
	private static final String START = "<S>";
	private static final String STOP = "</S>";
	private double vocabulary = 0.0;
	private double unseenCount = 0.0;
	private double seenCount = 0.0;
	
	public BigramModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.trainingCollection = trainingCollection;
		this.validationCollection = validationCollection;
		train();
		this.vocabulary = bigramCounterMap.keySet().size() - 1;
		this.seenCount = bigramCounterMap.totalSize();
		this.unseenCount = Math.pow(vocabulary, 2) - seenCount;
		System.out.println("vacabulary size: " + vocabulary);
		System.out.println("before smoothing, total seen type number is " + seenCount);
		System.out.println("before smoothing, total unseen type number is " + unseenCount);
		System.out.println("before smoothing, total seen token number is " + bigramCounterMap.totalCount());

		smooth();
		System.out.println("C[0] is " + C[0]);
		System.out.println("after smoothing, total token is " + validate() + ", and N[1] is " + N[1]);
		solveHistory();
	}
	
	void train() {
		for (List<String> sentence : trainingCollection) {
			List<String> stopSentence = new ArrayList<String>(sentence);
			stopSentence.add(0, START);
			stopSentence.add(STOP);
			
			Iterator<String> iter = stopSentence.iterator();
			String first = iter.next();
			String second = iter.next();

			while (iter.hasNext() || !first.equals(STOP)) {
				Counter<String> counter = bigramCounterMap.getCounter(first);
				counter.incrementCount(second, INCREMENT_FACTOR);
				first = second;
				if (iter.hasNext()) {
					second = iter.next();
				}
			}
		}
		
		// find count of counts
		for (String first : bigramCounterMap.keySet()) {
			Counter<String> counter = bigramCounterMap.getCounter(first);
			for (String second : counter.keySet()) {
				double count = counter.getCount(second);
				nCounter.incrementCount((int)count, INCREMENT_FACTOR);
			}
		}
	}
	
	void smooth() {
		Set<Integer> treeSet = new TreeSet<Integer>();
		treeSet.addAll(nCounter.keySet());
		Iterator<Integer> iter = treeSet.iterator();
		
		int next = 0;
		while (iter.hasNext() && next <= THRESHOLD) {
			next = iter.next();
			N[next] = nCounter.getCount(next);
		}
		N[0] = unseenCount;
		
		
		for(int c = 1; c <= THRESHOLD; c++) {
			double cStar = Math.log((c + 1) * N[c + 1] / N[c]
					- (c * (THRESHOLD + 1) * N[THRESHOLD + 1]) / N[1])
					- Math.log(1.0 - (THRESHOLD + 1) * N[THRESHOLD + 1] / N[1]);
			cStar = Math.pow(Math.E, cStar);
			C[c] = cStar;
		}
		C[0] = N[1] / unseenCount;
	}
	
	double validate() {
		double sum = 0.0;
		for (String first : bigramCounterMap.keySet()) {
			Counter<String> counter = bigramCounterMap.getCounter(first);
			for (String second : counter.keySet()) {
				double count = counter.getCount(second);
				if ((int)count <= THRESHOLD) {
					sum += C[(int)count];
				} else {
					sum += count;
				}
			}
		}
		return sum;
	}
	
	void solveHistory() {
		for (String first : bigramCounterMap.keySet()) {
			Counter<String> counter = bigramCounterMap.getCounter(first);
			double t = vocabulary + 1;
			double size = counter.size();
			double sum = 0.0;
			for (String second : counter.keySet()) {
				double count = counter.getCount(second);
				if ((int)count <= THRESHOLD) {
					sum += C[(int)count];
				} else {
					sum += count;
				}
			}
//			double marginal = (t - size) * C[0] + sum;
			double marginal = sum - 1.0 / (t - 1.0);
			historyCounter.incrementCount(first, marginal);
		}
	}
	
	
	double getWordLogProbability(List<String> sentence, int index) {
		if (index <= 0) {
			try {
				throw new Exception();
			} catch(Exception e) {
				System.err.println("Index out of range!");
			}
		}
		String history = sentence.get(index - 1);
		String word = sentence.get(index);
		
		// if history unseen
		return getPWord(history, word);
	}

	double getPWord(String history, String word) {
		if (!bigramCounterMap.containsKey(history)) {
			return Math.log(1.0 / (vocabulary + 1.0)) / Math.log(2.0);
		}
		
		double count = 0.0;
		Counter<String> counter = bigramCounterMap.getCounter(history);
		int intCount = (int) counter.getCount(word);
		if (intCount <= THRESHOLD) {
			count = C[intCount];
		} else {
			count = intCount;
		}
		double marginal = historyCounter.getCount(history);
		return Math.log(count / marginal) / Math.log(2.0);
	}
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		double probability = 0.0;
		for (int i = 1; i < stoppedSentence.size(); i++) {
			probability += getWordLogProbability(stoppedSentence, i);
		}
//		System.out.println("sentence p is " + probability);
		return Math.pow(2.0, probability);
	}

	String generateWord(String history) {
		double sample = Math.random();
		double sum = 0.0;
		Counter<String> counter = bigramCounterMap.getCounter(history);
		double marginal = historyCounter.getCount(history);
		for (String word : counter.keySet()) {
			int intCount = (int) counter.getCount(word);
			if (intCount <= THRESHOLD) {
				sum += C[intCount];
			} else {
				sum += intCount;
			}
			if (sum / marginal >= sample) {
				return word;
			}
		}
		return "*UNKNOWN*";
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		
		String word = generateWord(START);
		while (!word.equals(STOP)) {
			sentence.add(word);
			word = generateWord(word);
		}
		return sentence;
	}
	
	public CounterMap<String, String> getBigramCounterMap() {
		return this.bigramCounterMap;
	}
	/**
	 // test for probability sum add up to 1
	String generateWord(String history) {
		Counter<String> counter = bigramCounterMap.getCounter(history);
		double sum = 0.0;
		double marginal = historyCounter.getCount(history);
		for (String word : counter.keySet()) {
			int intCount = (int) counter.getCount(word);
			if (intCount <= THRESHOLD) {
				sum += C[intCount];
			} else {
				sum += intCount;
			}
		}
		System.out.println("sum/marginal + 1/t is " + (sum / marginal + 1.0 / (this.vocabulary + 1.0)));
		return null;
	}
	
	@Override
	public List<String> generateSentence() {
		// TODO Auto-generated method stub
		int i = 0;
		for (String first : bigramCounterMap.keySet()) {
			generateWord(first);
			i++;
			if (i > 10) {
				break;
			}
		}
		return null;
	}
	*/
}
